import torch
import torch.nn as nn


class Conv(nn.Module):
    def __init__(self, inchannel, outchannel, kernel=1):
        super().__init__()
        self.conv = nn.Conv2d(inchannel, outchannel, kernel_size=kernel, bias=False)
        self.bn = nn.BatchNorm2d(outchannel)
        self.relu = lambda x: x * torch.sigmoid(x)
        # self.relu = lambda x: x * torch.tanh(nn.functional.softplus(x))

    def forward(self, x):  # Nx768x4x5
        x = self.conv(x)  # Nx96x4x5
        x = self.bn(x)
        x = self.relu(x)
        return x


class ZhnNetFpn(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = Conv(192, 96)
        self.conv2 = Conv(192, 48)
        self.conv3 = Conv(96, 24)
        self.upsample = nn.Upsample(scale_factor=2)

    def forward(self, x):
        x3, x2, x1 = x
        y1 = self.conv1(x1)
        y2 = self.upsample(y1)
        y2 = torch.cat([x2, y2], dim=1)
        y2 = self.conv2(y2)
        y3 = self.upsample(y2)
        y3 = torch.cat([x3, y3], dim=1)
        y3 = self.conv3(y3)
        return y1, y2, y3
